{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 对CoH模块的测试"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 日期格式转化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 05-29 03:22:38 [__init__.py:239] Automatically detected platform cuda.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "'on the 2nd day'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from coh.utils import int_to_ordinal\n",
    "int_to_ordinal(48)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* LLM筛选文本处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['sda', '3', 'e']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from coh.utils import parse_ids_to_list\n",
    "parse_ids_to_list(\"1, 3, 5\", [\"asd\", \"sda\", \"a\", \"3\", \"t\", \"e\", \"w\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* vllm LLM模型加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING 05-29 03:22:40 [config.py:2972] Casting torch.bfloat16 to torch.float16.\n",
      "INFO 05-29 03:22:47 [config.py:717] This model supports multiple tasks: {'generate', 'embed', 'classify', 'score', 'reward'}. Defaulting to 'generate'.\n",
      "INFO 05-29 03:22:48 [gptq_marlin.py:147] Detected that the model can run with gptq_marlin, however you specified quantization=gptq explicitly, so forcing gptq. Use quantization=gptq_marlin for faster inference\n",
      "INFO 05-29 03:22:48 [gptq_bitblas.py:168] Detected that the model can run with gptq_bitblas, however you specified quantization=gptq explicitly, so forcing gptq. Use quantization=gptq_bitblas for faster inference\n",
      "WARNING 05-29 03:22:48 [config.py:830] gptq quantization is not fully optimized yet. The speed can be slower than non-quantized models.\n",
      "INFO 05-29 03:22:48 [config.py:2003] Chunked prefill is enabled with max_num_batched_tokens=8192.\n",
      "WARNING 05-29 03:22:48 [cuda.py:93] To see benefits of async output processing, enable CUDA graph. Since, enforce-eager is enabled, async output processor cannot be used\n",
      "WARNING 05-29 03:22:49 [utils.py:2382] We must use the `spawn` multiprocessing start method. Overriding VLLM_WORKER_MULTIPROC_METHOD to 'spawn'. See https://docs.vllm.ai/en/latest/getting_started/troubleshooting.html#python-multiprocessing for more information. Reason: CUDA is initialized\n",
      "INFO 05-29 03:22:52 [__init__.py:239] Automatically detected platform cuda.\n",
      "INFO 05-29 03:22:54 [core.py:58] Initializing a V1 LLM engine (v0.8.5.post1) with config: model='/root/autodl-fs/models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/', speculative_config=None, tokenizer='/root/autodl-fs/models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, override_neuron_config=None, tokenizer_revision=None, trust_remote_code=False, dtype=torch.float16, max_seq_len=32768, download_dir=None, load_format=auto, tensor_parallel_size=1, pipeline_parallel_size=1, disable_custom_all_reduce=False, quantization=gptq, enforce_eager=True, kv_cache_dtype=auto,  device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='auto', reasoning_backend=None), observability_config=ObservabilityConfig(show_hidden_metrics=False, otlp_traces_endpoint=None, collect_model_forward_time=False, collect_model_execute_time=False), seed=None, served_model_name=/root/autodl-fs/models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/, num_scheduler_steps=1, multi_step_stream_outputs=True, enable_prefix_caching=True, chunked_prefill_enabled=True, use_async_output_proc=False, disable_mm_preprocessor_cache=False, mm_processor_kwargs=None, pooler_config=None, compilation_config={\"splitting_ops\":[],\"compile_sizes\":[],\"cudagraph_capture_sizes\":[],\"max_capture_size\":0}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2025-05-29 03:22:55,195 - INFO - flashinfer.jit: Prebuilt kernels not found, using JIT backend\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING 05-29 03:22:55 [utils.py:2522] Methods determine_num_available_blocks,device_config,get_cache_block_size_bytes,initialize_cache not implemented in <vllm.v1.worker.gpu_worker.Worker object at 0x7fb720dad7c0>\n",
      "INFO 05-29 03:22:55 [parallel_state.py:1004] rank 0 in world size 1 is assigned as DP rank 0, PP rank 0, TP rank 0\n",
      "INFO 05-29 03:22:55 [cuda.py:221] Using Flash Attention backend on V1 engine.\n",
      "INFO 05-29 03:22:55 [topk_topp_sampler.py:44] Currently, FlashInfer top-p & top-k sampling sampler is disabled because FlashInfer>=v0.2.3 is not backward compatible. Falling back to the PyTorch-native implementation of top-p & top-k sampling.\n",
      "INFO 05-29 03:22:55 [gpu_model_runner.py:1329] Starting to load model /root/autodl-fs/models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/...\n",
      "WARNING 05-29 03:22:55 [utils.py:168] The model class MixtralForCausalLM has not defined `packed_modules_mapping`, this may lead to incorrect mapping of quantized or ignored modules\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Loading safetensors checkpoint shards:   0% Completed | 0/1 [00:00<?, ?it/s]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:02<00:00,  2.31s/it]\n",
      "Loading safetensors checkpoint shards: 100% Completed | 1/1 [00:02<00:00,  2.31s/it]\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 05-29 03:23:00 [loader.py:458] Loading weights took 2.80 seconds\n",
      "INFO 05-29 03:23:00 [gpu_model_runner.py:1347] Model loading took 22.1664 GiB and 4.799877 seconds\n",
      "INFO 05-29 03:23:06 [kv_cache_utils.py:634] GPU KV cache size: 38,048 tokens\n",
      "INFO 05-29 03:23:06 [kv_cache_utils.py:637] Maximum concurrency for 32,768 tokens per request: 1.16x\n",
      "INFO 05-29 03:23:06 [core.py:159] init engine (profile, create kv cache, warmup model) took 5.53 seconds\n",
      "INFO 05-29 03:23:06 [core_client.py:439] Core engine process 0 ready.\n"
     ]
    }
   ],
   "source": [
    "from vllm import LLM\n",
    "import torch\n",
    "\n",
    "model_name = \"mixtral\"\n",
    "model_path = \"/root/autodl-fs/models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/\"\n",
    "\n",
    "num_gpus = 1\n",
    "\n",
    "llm = LLM(\n",
    "    model=model_path,\n",
    "    quantization=\"gptq\",\n",
    "    dtype=torch.float16,\n",
    "    tensor_parallel_size=num_gpus,\n",
    "    enforce_eager=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* transformers LLM模型加载测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "# import torch\n",
    "\n",
    "# model_name = \"./models/TheBloke/Mixtral-8x7B-Instruct-v0___1-GPTQ/\"\n",
    "\n",
    "# tokenizer = AutoTokenizer.from_pretrained(\n",
    "#     model_name, \n",
    "#     trust_remote_code=True,\n",
    "#     )\n",
    "# tokenizer.pad_token = tokenizer.eos_token\n",
    "# llm = AutoModelForCausalLM.from_pretrained(\n",
    "#     model_name, \n",
    "#     trust_remote_code=True,\n",
    "#     device_map=\"auto\",\n",
    "#     torch_dtype=torch.float16\n",
    "#     )\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 数据集加载测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cache for fliter test data found.\n",
      "seen dataset proportion: 0.9345787324156708\n",
      "unseen dataset proportion: 0.06542126758432915\n"
     ]
    }
   ],
   "source": [
    "from utils import Data\n",
    "\n",
    "params_fliter = {\n",
    "    \"max_tokens\": 16,\n",
    "    \"top_p\": 1.0,\n",
    "    \"temperature\": 0.0,\n",
    "    \"stop\": [\"</s>\", \"[/INST]\"]\n",
    "    }\n",
    "\n",
    "contents = Data(\n",
    "    model_name=model_name,\n",
    "    llm = llm,\n",
    "    param = params_fliter,\n",
    "    \n",
    "    dataset=\"ICEWS14_forecasting\", \n",
    "    add_reverse_relation=True\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* LLM预测文本处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['independent corrupt practices commission', 'education (nigeria)', 'muslim (nigeria)', 'member of the judiciary (nigeria)', 'citizen (nigeria)', 'boko haram']\n"
     ]
    }
   ],
   "source": [
    "from coh.utils import parse_predict_answer\n",
    "\n",
    "_str = \"\"\"Possible answers:\n",
    "1. Independent Corrupt Practices Commission\n",
    "2. Education (Nigeria)\n",
    "3. Muslim (Nigeria)\n",
    "4. Member of the Judiciary (Nigeria)\n",
    "5. Citizen (Nigeria)\n",
    "6. Boko Haram\n",
    "7. Government (Nigeria)\n",
    "\"\"\"\n",
    "\n",
    "print(parse_predict_answer(_str, contents.entity2id))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* CoH模块加载测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-29 03:23:07] INFO model.py:58: testing LLM...\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "824b5c160a27438ba0d8dfa132cbe4ff",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processed prompts:   0%|          | 0/1 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[2025-05-29 03:23:15] INFO model.py:60: LLM is working!\n"
     ]
    }
   ],
   "source": [
    "from coh import CoH\n",
    "\n",
    "params_coh = {\n",
    "    \"max_tokens\": 8000,\n",
    "    \"top_p\": 1.0,\n",
    "    \"temperature\": 0.0,\n",
    "    \"stop\": [\"</s>\", \"[/INST]\"]\n",
    "    }\n",
    "\n",
    "coh = CoH(\n",
    "    num_entities=contents.num_entities,\n",
    "    num_relations=contents.num_relations,\n",
    "    entity2id=contents.entity2id,\n",
    "    id2entity=contents.id2entity,\n",
    "    id2relation=contents.id2relation,\n",
    "    s_his_dict=contents.get_adj_his_format_dict(),\n",
    "    llm=llm,\n",
    "    # tokenizer=tokenizer,\n",
    "    params=params_coh,\n",
    "    expand_n=5\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 取10条数据测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([np.int64(0), np.int64(18), np.int64(7536)], np.int64(384))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_data = contents.test_data\n",
    "test_data = test_data[:10]\n",
    "\n",
    "def extract_data(data):\n",
    "    s, p, o, t = data\n",
    "    return [s, p, t], o\n",
    "\n",
    "q, a = zip(*[extract_data(item) for item in test_data])\n",
    "\n",
    "q[0], a[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 推理过程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e28a0e894c264d27adfa3de997cdf8ce",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processed prompts:   0%|          | 0/10 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "077c9b68a4884c099164c412b8171fc7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processed prompts:   0%|          | 0/10 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e5e103fff98d43b7ba660c015bae64b7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Processed prompts:   0%|          | 0/10 [00:00<?, ?it/s, est. speed input: 0.00 toks/s, output: 0.00 toks/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "scores = coh(q)[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 样例研究"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      "target: Cabinet / Council of Ministers / Advisors (United States)\tConsult\tSergey Viktorovich Lavrov\ton the 314th day\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Benjamin Netanyahu\tConsult\tLikud\ton the 314th day\n",
      "\n",
      "UN Security Council: tensor(0.4256)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Japan\tEngage in diplomatic cooperation\tChina\ton the 314th day\n",
      "\n",
      "South Korea: tensor(0.3543)\n",
      "\n",
      "China: tensor(0.4256)\n",
      "\n",
      "France: tensor(0.2891)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Japan\tArrest, detain, or charge with legal action\tCitizen (South Korea)\ton the 314th day\n",
      "\n",
      "Kim Jong-Un: tensor(0.4256)\n",
      "\n",
      "Xi Jinping: tensor(0.3543)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Japan\tMake an appeal or request\tChina\ton the 314th day\n",
      "\n",
      "China: tensor(0.4256)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Japan\tExpress intent to engage in diplomatic cooperation (such as policy support)\tAssociation of Southeast Asian Nations\ton the 314th day\n",
      "\n",
      "North Korea: tensor(0.3543)\n",
      "\n",
      "South Korea: tensor(0.2891)\n",
      "\n",
      "China: tensor(0.4256)\n",
      "\n",
      "France: tensor(0.2315)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: Japan\tEngage in negotiation\tSouth Korea\ton the 314th day\n",
      "\n",
      "North Korea: tensor(0.3543)\n",
      "\n",
      "South Korea: tensor(0.2891)\n",
      "\n",
      "China: tensor(0.4256)\n",
      "\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------\n",
      "target: Japan\tExpress intent to cooperate\tChina\ton the 314th day\n",
      "\n",
      "China: tensor(0.4256)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: North Korea\tReduce relations\tCitizen (North Korea)\ton the 314th day\n",
      "\n",
      "Japan: tensor(0.3543)\n",
      "\n",
      "South Korea: tensor(0.4256)\n",
      "\n",
      "China: tensor(0.2891)\n",
      "\n",
      "--------------------------------------------------\n",
      "target: North Korea\tCriticize or denounce\tSouth Korea\ton the 314th day\n",
      "\n",
      "Japan: tensor(0.2891)\n",
      "\n",
      "South Korea: tensor(0.4256)\n",
      "\n",
      "China: tensor(0.3543)\n",
      "\n",
      "Citizen (South Korea): tensor(0.1824)\n",
      "\n",
      "Head of Government (South Korea): tensor(0.2315)\n",
      "\n",
      "Detainee (United States): tensor(0.1419)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "scores.shape\n",
    "scores # (batch_size, num_entities)\n",
    "\n",
    "q_list = coh.format_his_list([s, p, o, t] for [s, p, t], o in zip(q, a))\n",
    "\n",
    "for score, _q in zip(scores, q_list):\n",
    "    print(\"--------------------------------------------------\")\n",
    "    print(\"target: \" + _q + '\\n')\n",
    "    for i in range(len(score)):\n",
    "        if score[i] != 0:\n",
    "            print(contents.id2entity[i] + ': ' + str(score[i]))\n",
    "            print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* 评估测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--- coh test ---\n",
      "MRR (%): 44.40\n",
      "Hits@1 (%): 40.00\n",
      "Hits@3 (%): 50.00\n",
      "Hits@10 (%): 60.00\n",
      "-------------------------\n"
     ]
    }
   ],
   "source": [
    "from utils import Evaluator\n",
    "\n",
    "eval = Evaluator(\"coh test\")\n",
    "eval.update(score_batch=scores, true_list=a)\n",
    "\n",
    "eval.print()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.21"
  },
  "toc-autonumbering": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
